import streamlit as st
import cv2
from services import db, face_recog

st.title("录入员工信息 & 采集样本")

db.init_db()

# 录入员工信息
with st.form("add_employee_form"):
    emp_id = st.text_input("工号")
    name = st.text_input("姓名")
    dept = st.text_input("部门")
    note = st.text_area("备注")
    submitted = st.form_submit_button("保存信息")
    if submitted:
        if emp_id and name and dept:
            ok,msg = db.add_employee(emp_id,name,dept,note)
            if ok:
                st.success(msg)
            else:
                st.error(msg)
        else:
            st.warning("请完整填写工号、姓名、部门！")

st.markdown("---")
st.subheader("采集人脸样本")

# 摄像头预览
frame_slot = st.empty()

def camera_preview(emp_id, save=False, manual=False):
    """摄像头预览 + 采集样本"""
    global num
    cap = cv2.VideoCapture(0)
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
    count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)
        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x,y), (x+w,y+h), (0,255,0), 2)
            if save:
                face_img = gray[y:y+h, x:x+w]
                filepath, num = face_recog.save_sample(emp_id, face_img)
                count += 1
                save = False if manual else save  # 手动只采一次
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame_slot.image(frame_rgb, channels="RGB")
        if manual or count >= 30:  # 手动只采 1 张 / 自动采 30 张
            break
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()
    return count, num

# 自动采样
with st.form("auto_sample_form"):
    emp_id_auto = st.text_input("输入工号自动采样", key="auto_emp_id")
    auto_submit = st.form_submit_button("开始自动采样（30张）")
    if auto_submit:
        if emp_id_auto:
            count, num = camera_preview(emp_id_auto, save=True, manual=False)
            st.success(f"自动采集完成，共保存 {count} 张样本, 当前样本数量: {num}")
        else:
            st.warning("请先输入工号")

# 手动采样
with st.form("manual_sample_form"):
    emp_id_manual = st.text_input("输入工号手动采样", key="manual_emp_id")
    manual_submit = st.form_submit_button("采集一张样本")
    if manual_submit:
        if emp_id_manual:
            count, num = camera_preview(emp_id_manual, save=True, manual=True)
            if count > 0:
                st.success(f"已保存 1 张样本, 当前样本数量: {num}")
            else:
                st.warning("未检测到人脸，请重试")
        else:
            st.warning("请先输入工号")

# 查看样本
with st.form("preview_sample_form"):
    emp_id_preview = st.text_input("输入工号查看样本", key="preview_emp_id")
    preview_submit = st.form_submit_button("查看样本")
    if preview_submit:
        if emp_id_preview:
            imgs = face_recog.preview_samples(emp_id_preview)
            if imgs:
                captions = [f"样本{i+1}" for i in range(len(imgs))]
                st.image(imgs, caption=captions, width=150)
            else:
                st.warning("未找到样本，请先采集")
        else:
            st.warning("请先输入工号")

# 训练模型
with st.form("train_model_form"):
    train_submit = st.form_submit_button("训练模型")
    if train_submit:
        ok,msg = face_recog.train_and_save_model()
        if ok:
            st.success(msg)
        else:
            st.error(msg)
